#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <time.h>
#include <signal.h>
#include <errno.h>
#include <pthread.h>
#include <sys/time.h>
#ifdef HAVE_GETOPT_H
#include <getopt.h>
#else
// 简单的getopt实现
extern char *optarg;
extern int optind, opterr, optopt;
int getopt(int argc, char * const argv[], const char *optstring);
#endif

#define BUFFER_SIZE 1024 * 1024  // 1MB buffer
#define DEFAULT_PORT 8888
#define DEFAULT_SIZE 100  // 默认发送100MB
#define DEFAULT_THREADS 4  // 默认4个并发连接

// 全局变量
volatile int running = 1;
pthread_mutex_t stats_mutex = PTHREAD_MUTEX_INITIALIZER;
long total_bytes = 0;
time_t start_time, end_time;

// 线程参数结构
typedef struct {
    int thread_id;
    char *server_ip;
    int port;
    long target_bytes;
    int test_time_sec;
    int use_time_mode;
    long thread_bytes;
    int client_fd;
} thread_args_t;

void signal_handler(int sig) {
    (void)sig;  // 避免未使用参数警告
    printf("\n正在停止传输...\n");
    running = 0;
}

void print_usage(const char *program_name) {
    printf("用法: %s [选项]\n", program_name);
    printf("选项:\n");
    printf("  -s <服务器IP>  指定服务器IP地址 (默认: 127.0.0.1)\n");
    printf("  -p <端口>      指定服务器端口 (默认: %d)\n", DEFAULT_PORT);
    printf("  -d <大小>      指定发送数据大小(MB) (默认: %d)\n", DEFAULT_SIZE);
    printf("  -t <时间>      指定测试时间(秒) (与-d互斥)\n");
    printf("  -n <线程数>    指定并发连接数 (默认: %d)\n", DEFAULT_THREADS);
    printf("  -h             显示此帮助信息\n");
    printf("\n示例:\n");
    printf("  %s -s 192.168.1.100 -d 500 -n 8    # 8个并发连接发送500MB数据\n", program_name);
    printf("  %s -s 192.168.1.100 -t 60 -n 16    # 16个并发连接发送60秒数据\n", program_name);
    printf("\n注意: 多线程测试可以充分利用bond网络的聚合带宽\n");
}

// 工作线程函数
void* worker_thread(void* arg) {
    thread_args_t *args = (thread_args_t*)arg;
    char buffer[BUFFER_SIZE];
    ssize_t bytes_sent;
    long thread_total = 0;
    
    // 初始化发送缓冲区
    memset(buffer, 'A', BUFFER_SIZE);
    
    // 创建socket
    args->client_fd = socket(AF_INET, SOCK_STREAM, 0);
    if (args->client_fd == -1) {
        perror("创建socket失败");
        return NULL;
    }
    
    // 配置服务器地址
    struct sockaddr_in server_addr;
    memset(&server_addr, 0, sizeof(server_addr));
    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(args->port);
    
    if (inet_pton(AF_INET, args->server_ip, &server_addr.sin_addr) <= 0) {
        fprintf(stderr, "线程%d: 无效的IP地址: %s\n", args->thread_id, args->server_ip);
        close(args->client_fd);
        return NULL;
    }
    
    // 连接到服务器
    if (connect(args->client_fd, (struct sockaddr *)&server_addr, sizeof(server_addr)) == -1) {
        perror("连接服务器失败");
        close(args->client_fd);
        return NULL;
    }
    
    printf("线程%d: 连接成功\n", args->thread_id);
    
    // 开始传输数据
    while (running) {
        // 检查是否达到目标
        if (!args->use_time_mode && thread_total >= args->target_bytes) {
            break;
        }
        if (args->use_time_mode && (time(NULL) - start_time) >= args->test_time_sec) {
            break;
        }
        
        // 发送数据
        bytes_sent = send(args->client_fd, buffer, BUFFER_SIZE, 0);
        if (bytes_sent == -1) {
            if (errno == EINTR) {
                continue;  // 被信号中断，继续循环
            }
            perror("发送数据失败");
            break;
        }
        
        thread_total += bytes_sent;
        
        // 更新全局统计
        pthread_mutex_lock(&stats_mutex);
        total_bytes += bytes_sent;
        pthread_mutex_unlock(&stats_mutex);
    }
    
    args->thread_bytes = thread_total;
    close(args->client_fd);
    
    printf("线程%d: 传输完成，发送 %.2f MB\n", 
           args->thread_id, thread_total / (1024.0 * 1024.0));
    
    return NULL;
}

// 统计显示线程
void* stats_thread(void* arg) {
    thread_args_t *args = (thread_args_t*)arg;
    time_t last_update_time = start_time;
    
    while (running) {
        sleep(1);
        
        time_t current_time = time(NULL);
        if (current_time != last_update_time) {
            double duration = difftime(current_time, start_time);
            if (duration > 0) {
                pthread_mutex_lock(&stats_mutex);
                long current_total = total_bytes;
                pthread_mutex_unlock(&stats_mutex);
                
                double bandwidth_mbps = (current_total * 8.0) / (duration * 1024 * 1024);
                
                if (args->use_time_mode) {
                    printf("\r总发送: %.2f MB, 聚合带宽: %.2f Mbps, 剩余时间: %ld 秒", 
                           current_total / (1024.0 * 1024.0), bandwidth_mbps, 
                           args->test_time_sec - (current_time - start_time));
                } else {
                    printf("\r总发送: %.2f MB, 聚合带宽: %.2f Mbps, 进度: %.1f%%", 
                           current_total / (1024.0 * 1024.0), bandwidth_mbps,
                           (double)current_total / (args->target_bytes * args->thread_id) * 100);
                }
                fflush(stdout);
            }
            last_update_time = current_time;
        }
    }
    
    return NULL;
}

int main(int argc, char *argv[]) {
    char *server_ip = "127.0.0.1";
    int port = DEFAULT_PORT;
    int data_size_mb = DEFAULT_SIZE;
    int test_time_sec = 0;
    int num_threads = DEFAULT_THREADS;
    int opt;
    int use_time_mode = 0;
    pthread_t *threads;
    pthread_t stats_tid;
    thread_args_t *thread_args;
    int i;
    double duration, bandwidth_mbps;
    
    // 解析命令行参数
    while ((opt = getopt(argc, argv, "s:p:d:t:n:h")) != -1) {
        switch (opt) {
            case 's':
                server_ip = optarg;
                break;
            case 'p':
                port = atoi(optarg);
                if (port <= 0 || port > 65535) {
                    fprintf(stderr, "错误: 端口号必须在1-65535之间\n");
                    exit(EXIT_FAILURE);
                }
                break;
            case 'd':
                data_size_mb = atoi(optarg);
                if (data_size_mb <= 0) {
                    fprintf(stderr, "错误: 数据大小必须大于0\n");
                    exit(EXIT_FAILURE);
                }
                break;
            case 't':
                test_time_sec = atoi(optarg);
                if (test_time_sec <= 0) {
                    fprintf(stderr, "错误: 测试时间必须大于0\n");
                    exit(EXIT_FAILURE);
                }
                use_time_mode = 1;
                break;
            case 'n':
                num_threads = atoi(optarg);
                if (num_threads <= 0 || num_threads > 64) {
                    fprintf(stderr, "错误: 线程数必须在1-64之间\n");
                    exit(EXIT_FAILURE);
                }
                break;
            case 'h':
                print_usage(argv[0]);
                exit(EXIT_SUCCESS);
            default:
                print_usage(argv[0]);
                exit(EXIT_FAILURE);
        }
    }
    
    // 检查参数冲突
    if (use_time_mode && data_size_mb != DEFAULT_SIZE) {
        fprintf(stderr, "错误: -d和-t参数不能同时使用\n");
        exit(EXIT_FAILURE);
    }
    
    // 设置信号处理
    signal(SIGINT, signal_handler);
    signal(SIGTERM, signal_handler);
    
    // 分配内存
    threads = malloc(num_threads * sizeof(pthread_t));
    thread_args = malloc(num_threads * sizeof(thread_args_t));
    if (!threads || !thread_args) {
        fprintf(stderr, "错误: 内存分配失败\n");
        exit(EXIT_FAILURE);
    }
    
    printf("多线程TCP带宽测试客户端\n");
    printf("========================\n");
    printf("服务器: %s:%d\n", server_ip, port);
    printf("并发连接数: %d\n", num_threads);
    if (use_time_mode) {
        printf("测试时间: %d 秒\n", test_time_sec);
    } else {
        printf("总数据大小: %d MB (每个连接: %.2f MB)\n", 
               data_size_mb, (double)data_size_mb / num_threads);
    }
    printf("按Ctrl+C停止传输\n\n");
    
    // 初始化线程参数
    for (i = 0; i < num_threads; i++) {
        thread_args[i].thread_id = i + 1;
        thread_args[i].server_ip = server_ip;
        thread_args[i].port = port;
        thread_args[i].use_time_mode = use_time_mode;
        thread_args[i].test_time_sec = test_time_sec;
        thread_args[i].target_bytes = use_time_mode ? 0 : (long)data_size_mb * 1024 * 1024 / num_threads;
        thread_args[i].thread_bytes = 0;
        thread_args[i].client_fd = -1;
    }
    
    // 记录开始时间
    start_time = time(NULL);
    
    // 创建统计显示线程
    if (pthread_create(&stats_tid, NULL, stats_thread, &thread_args[0]) != 0) {
        perror("创建统计线程失败");
        free(threads);
        free(thread_args);
        exit(EXIT_FAILURE);
    }
    
    // 创建工作线程
    for (i = 0; i < num_threads; i++) {
        if (pthread_create(&threads[i], NULL, worker_thread, &thread_args[i]) != 0) {
            perror("创建工作线程失败");
            running = 0;
            break;
        }
    }
    
    // 等待所有工作线程完成
    for (i = 0; i < num_threads; i++) {
        pthread_join(threads[i], NULL);
    }
    
    // 停止统计线程
    running = 0;
    pthread_join(stats_tid, NULL);
    
    // 计算最终统计信息
    end_time = time(NULL);
    duration = difftime(end_time, start_time);
    
    printf("\n\n=== 聚合带宽测试结果 ===\n");
    printf("并发连接数: %d\n", num_threads);
    printf("总发送数据: %.2f MB\n", total_bytes / (1024.0 * 1024.0));
    printf("传输时间: %.2f 秒\n", duration);
    if (duration > 0) {
        bandwidth_mbps = (total_bytes * 8.0) / (duration * 1024 * 1024);
        printf("聚合带宽: %.2f Mbps\n", bandwidth_mbps);
        printf("平均速度: %.2f MB/s\n", total_bytes / (duration * 1024 * 1024));
        printf("单连接平均带宽: %.2f Mbps\n", bandwidth_mbps / num_threads);
    }
    
    // 显示每个线程的统计
    printf("\n各线程统计:\n");
    for (i = 0; i < num_threads; i++) {
        printf("  线程%d: %.2f MB\n", i + 1, thread_args[i].thread_bytes / (1024.0 * 1024.0));
    }
    printf("========================\n");
    
    // 清理资源
    free(threads);
    free(thread_args);
    pthread_mutex_destroy(&stats_mutex);
    
    return 0;
}
